from osgeo import ogr
import numpy as np
import matplotlib.pyplot as plt
import geopandas as gpd

def read_shapefile(shapefile_path):
    driver = ogr.GetDriverByName('ESRI Shapefile')
    dataset = driver.Open(shapefile_path)
    layer = dataset.GetLayer()
    gdf = gpd.read_file(shapefile_path)
    
    # 获取字段名
    field_names = [field.GetName() for field in layer.schema]

    # 获取要素与属性
    geometries = []
    attributes = []

    for feature in layer:
        geometry = feature.GetGeometryRef()
        attributes.append([feature.GetField(field) for field in field_names])
        geometries.append(geometry)

    dataset = None 
    return geometries, field_names, attributes,gdf

def plot_map(gdf):
    # 绘制地图
    gdf.plot()

    # 显示地图
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.title('Vector Map')
    plt.show()


def display_attribute_table(field_names, attributes):
    print("{:<10}".format("Index"), end="")
    ##输出字段名
    for name in field_names:
        print("{:<15}".format(name), end="")
    print()  
    ##输出字段值
    for i, row in enumerate(attributes):
        print("{:<10}".format(i), end="")
        for value in row:
            print("{:<15}".format(value), end="")
        print() 

def write_shapefile(output_path, geometries, field_names, attributes):
    driver = ogr.GetDriverByName('ESRI Shapefile')

    if driver is None:
        raise Exception("ESRI Shapefile driver not available")

    # 创建一个dataset
    dataset = driver.CreateDataSource(output_path)
    layer = dataset.CreateLayer('layer', geom_type=ogr.wkbPolygon)

    # 创建字段
    for field_name in field_names:
        field_defn = ogr.FieldDefn(field_name, ogr.OFTString)
        layer.CreateField(field_defn)

    # 写入要素几何与字段
    for i, geometry in enumerate(geometries):
        feature = ogr.Feature(layer.GetLayerDefn())
        feature.SetGeometry(geometry)

        for j, value in enumerate(attributes[i]):
            feature.SetField(field_names[j], str(value))

        layer.CreateFeature(feature)
        feature = None  # dereference the feature

    dataset = None  # close the dataset




def write_shapefile1(source_path,output_path):
    ds = ogr.Open(source_path, 0)           # 0是只读，1是可写
    layer = ds.GetLayer()    # 获取矢量文件Layer
    ### 
    driver = ogr.GetDriverByName('ESRI Shapefile')
    ds_out = driver.CreateDataSource(output_path)
    layer_out = ds_out.CreateLayer('vector_new_2', geom_type=ogr.wkbPolygon, srs=layer.GetSpatialRef())
    fea_defn = layer_out.GetLayerDefn()    ## 获得空要素定义(或空图层定义)

    ### 创建字段
    layer_defn = layer_out.GetLayerDefn()
    print('Number of fields (before field creating): ', layer_defn.GetFieldCount())  ### 字段创建前字段数
    fea = layer.GetFeature(0)
    for i in range(fea.GetFieldCount()):
        field_defn = fea.GetFieldDefnRef(i)
        layer_out.CreateField(field_defn)
    layer_defn = layer_out.GetLayerDefn()
    print('Number of fields (after field creating): ', layer_defn.GetFieldCount())

    for fea in layer:
        fea_out = ogr.Feature(fea_defn)
        fea_out.SetGeometry(fea.geometry())
        for i in range(fea.GetFieldCount()):
            field_value = fea.GetField(i)
            fea_out.SetField(i, field_value)
        layer_out.CreateFeature(fea_out)
    ds_out = None        
    print("done")